#include <sys/stat.h>
#include <fstream>
#include <glog/logging.h>

#include "SampleDetector.hpp"
#include "opencv2/imgproc.hpp"
#include "opencv2/imgcodecs.hpp"

#define TRY_RELEASE_POINTER(val, func) \
    if (val != nullptr)                \
    {                                  \
        func(val);                     \
        val = nullptr;                 \
    }
#define TRY_DELETE_POINTER(val, func) \
    if (val != nullptr)               \
    {                                 \
        func val;                     \
        val = nullptr;                \
    }

static bool ifFileExists(const char *FileName)
{
    struct stat my_stat
    {
    };
    return (stat(FileName, &my_stat) == 0);
}

static bool replace(std::string &str, const std::string &from, const std::string &to)
{
    size_t start_pos = str.find(from);
    if (start_pos == std::string::npos)
        return false;
    str.replace(start_pos, from.length(), to);
    return true;
}

static float IOU(const cv::Rect &b1, const cv::Rect &b2)
{
    auto intersec = b1 & b2;
    return static_cast<float>(intersec.area()) / (b1.area() + b2.area() - intersec.area());
}

SampleDetector::SampleDetector()
{
}

STATUS SampleDetector::Init(const std::string &modelpath, const float threshold)
{
    threshold_ = threshold;
    
    std::string strModelName = modelpath;
    replace(strModelName, ".onnx", ".om");

    size_t sep_pos = strModelName.find_last_of(".");
    std::string om_Name = strModelName.substr(0, sep_pos);
    std::string soc_version = aclrtGetSocName();
    SDKLOG(INFO) << "soc_version : " << soc_version;
    if (!ifFileExists(strModelName.c_str()))
    {
        std::string cmd_str = "atc --framework=5 --model=" + modelpath + " --output=" + om_Name + "  --soc_version=" + soc_version;
        SDKLOG(INFO) << "system cmd status: " << system(cmd_str.c_str());
    }

    SDKLOG(INFO) << "start to init model ";
    m_acl_context = nullptr;
    m_acl_stream = nullptr;
    uint32_t deviceCount;
    auto ret = aclrtGetDeviceCount(&deviceCount);
    if (ret != ACL_ERROR_NONE)
    {
        SDKLOG(ERROR) << "No device found! aclError= " << ret;
        return ERROR_INITACL;
    }
    SDKLOG(INFO) << deviceCount << " devices found";

    ret = aclrtSetDevice(m_model_id);
    if (ret != ACL_ERROR_NONE)
    {
        SDKLOG(ERROR) << "Acl open device " << m_model_id << " failed! aclError= " << ret;
        return ERROR_INITACL;
    }
    SDKLOG(INFO) << "Open device " << m_model_id << " success";

    ret = aclrtCreateContext(&m_acl_context, m_model_id);
    if (ret != ACL_ERROR_NONE)
    {
        SDKLOG(ERROR) << "acl create context failed! aclError= " << ret;
        return ERROR_INITACL;
    }
    SDKLOG(INFO) << "create context success";

    ret = aclrtSetCurrentContext(m_acl_context);
    if (ret != ACL_ERROR_NONE)
    {
        SDKLOG(ERROR) << "acl set context failed! aclError= " << ret;
        return ERROR_INITACL;
    }
    SDKLOG(INFO) << "set context success";

    ret = aclrtCreateStream(&m_acl_stream);
    if (ret != 0)
    {
        SDKLOG(ERROR) << "failed to create  stream! aclError= " << ret;
        return ERROR_INITACL;
    }
    SDKLOG(INFO) << "create stream success";

    ret = aclrtGetRunMode(&m_acl_run_mode);
    if (ret != 0)
    {
        SDKLOG(ERROR) << "failed to get acl run mode! aclError= " << ret;
        return ERROR_INITACL;
    }
    if (m_acl_run_mode == ACL_DEVICE)
    {
        SDKLOG(INFO) << "run in device mode";
    }
    else
    {
        SDKLOG(INFO) << "run in host mode";
    }

    SDKLOG(INFO) << "load model " << strModelName;
    ret = aclmdlQuerySize(strModelName.c_str(), &mModelMSize, &mModelWSize);
    ret = aclrtMalloc(&mModelMptr, mModelMSize, ACL_MEM_MALLOC_HUGE_FIRST);
    ret = aclrtMalloc(&mModelWptr, mModelWSize, ACL_MEM_MALLOC_HUGE_FIRST);
    ret = aclmdlLoadFromFileWithMem(strModelName.c_str(), &m_model_id, mModelMptr, mModelMSize, mModelWptr, mModelWSize);
    m_model_desc = aclmdlCreateDesc();
    ret = aclmdlGetDesc(m_model_desc, m_model_id);

    int input_num = aclmdlGetNumInputs(m_model_desc);
    if (input_num != 1)
    {
        SDKLOG(ERROR) << "this is not a standard yolov8 model";
    }

    aclmdlGetInputDims(m_model_desc, 0, &m_input_dims);
    if (m_input_dims.dimCount != 4)
    {
        SDKLOG(ERROR) << "this is not a standard yolov8 model";
    }
    SDKLOG(INFO) << "input dim is : " << m_input_dims.dims[0] << " " << m_input_dims.dims[1] << " " << m_input_dims.dims[2] << " " << m_input_dims.dims[3];
    m_InputSize = cv::Size(m_input_dims.dims[3], m_input_dims.dims[2]);

    int output_num = aclmdlGetNumOutputs(m_model_desc);
    if (output_num != 1)
    {
        SDKLOG(ERROR) << "this is not a standard yolov8 model";
    }

    aclmdlGetOutputDims(m_model_desc, 0, &m_output_dims);
    if (m_output_dims.dimCount != 3)
    {
        SDKLOG(ERROR) << "this is not a standard yolov8 model";
    }
    SDKLOG(INFO) << "output dim is : " << m_output_dims.dims[0] << " " << m_output_dims.dims[1] << " " << m_output_dims.dims[2];
    m_iClassNums = m_output_dims.dims[1] - 4;
    m_iBoxNums = m_output_dims.dims[2];

    m_input_dataset = aclmdlCreateDataset();
    ret = aclrtMalloc(&m_input_buffer, aclmdlGetInputSizeByIndex(m_model_desc, 0), ACL_MEM_MALLOC_NORMAL_ONLY);
    m_input_data_buffer = aclCreateDataBuffer(m_input_buffer, aclmdlGetInputSizeByIndex(m_model_desc, 0));
    aclmdlAddDatasetBuffer(m_input_dataset, m_input_data_buffer);

    if (m_acl_run_mode == ACL_DEVICE)
    {
        SDKLOG(INFO) << "atlas run in device mode";
        size_t single_chn_size = m_input_dims.dims[0] * m_input_dims.dims[2] * m_input_dims.dims[3] * sizeof(float);
        m_chw_wrappers.emplace_back(m_input_dims.dims[2], m_input_dims.dims[3], CV_32FC1, (char *)m_input_buffer + 0 * single_chn_size);
        m_chw_wrappers.emplace_back(m_input_dims.dims[2], m_input_dims.dims[3], CV_32FC1, (char *)m_input_buffer + 1 * single_chn_size);
        m_chw_wrappers.emplace_back(m_input_dims.dims[2], m_input_dims.dims[3], CV_32FC1, (char *)m_input_buffer + 2 * single_chn_size);
    }

    m_output_dataset = aclmdlCreateDataset();
    ret = aclrtMalloc(&m_output_buffer, aclmdlGetOutputSizeByIndex(m_model_desc, 0), ACL_MEM_MALLOC_NORMAL_ONLY);
    if (m_acl_run_mode == ACL_HOST)
        m_output_buffer_host = new char[aclmdlGetOutputSizeByIndex(m_model_desc, 0)]();
    m_output_data_buffer = aclCreateDataBuffer(m_output_buffer, aclmdlGetOutputSizeByIndex(m_model_desc, 0));
    aclmdlAddDatasetBuffer(m_output_dataset, m_output_data_buffer);

    m_init_flag = true;
    return STATUS_SUCCESS;
}

STATUS SampleDetector::UnInit()
{
    if (m_init_flag == false)
    {
        return false;
    }
    SDKLOG(INFO) << "in uninit func";
    TRY_RELEASE_POINTER(m_input_dataset, aclmdlDestroyDataset);
    TRY_RELEASE_POINTER(m_input_data_buffer, aclDestroyDataBuffer);
    TRY_RELEASE_POINTER(m_input_buffer, aclrtFree);

    TRY_RELEASE_POINTER(m_output_dataset, aclmdlDestroyDataset);
    TRY_RELEASE_POINTER(m_output_data_buffer, aclDestroyDataBuffer);
    TRY_RELEASE_POINTER(mModelMptr, aclrtFree);
    TRY_RELEASE_POINTER(mModelWptr, aclrtFree);

    TRY_RELEASE_POINTER(m_output_buffer_host, free);
    TRY_RELEASE_POINTER(m_output_buffer, aclrtFree);

    aclmdlDestroyDesc(m_model_desc);
    if (m_model_id > 0)
    {
        aclmdlUnload(m_model_id);
    }
    aclrtDestroyStream(m_acl_stream);
    aclrtDestroyContext(m_acl_context);
    m_chw_wrappers.clear();
    m_init_flag = false;
}

SampleDetector::~SampleDetector()
{
    UnInit();
}

STATUS SampleDetector::ProcessImage(const cv::Mat &img, std::vector<BoxInfo> &DetObjs, const float threshold)
{
    auto ret = aclrtSetCurrentContext(m_acl_context);
    if (ret != ACL_ERROR_NONE)
    {
        SDKLOG(ERROR) << "acl set context failed! aclError= " << ret;
        return ERROR_INITACL;
    }

    threshold_ = threshold;
    DetObjs.clear();
    float r = std::min((float)m_InputSize.height / static_cast<float>(img.rows), (float)m_InputSize.width / static_cast<float>(img.cols));
    cv::Size new_size = cv::Size{int((float)img.cols * r), int((float)img.rows * r)};
    cv::Mat tmp_resized;

    cv::resize(img, tmp_resized, new_size);
    cv::cvtColor(tmp_resized, tmp_resized, cv::COLOR_BGR2RGB);
    m_Resized = cv::Mat(cv::Size(m_InputSize.width, m_InputSize.height), CV_8UC3, cv::Scalar(114, 114, 114));
    tmp_resized.copyTo(m_Resized(cv::Rect{0, 0, tmp_resized.cols, tmp_resized.rows}));

    m_Resized.convertTo(m_Normalized, CV_32FC3, 1 / 255.);
    cv::split(m_Normalized, m_chw_wrappers);

    if (m_acl_run_mode == ACL_HOST)
    {
        size_t single_chn_size = m_chw_wrappers[0].rows * m_chw_wrappers[0].cols * sizeof(float);
        aclrtMemcpy((char *)m_input_buffer + 0 * single_chn_size, single_chn_size, m_chw_wrappers[0].data, single_chn_size, ACL_MEMCPY_HOST_TO_DEVICE);
        aclrtMemcpy((char *)m_input_buffer + 1 * single_chn_size, single_chn_size, m_chw_wrappers[1].data, single_chn_size, ACL_MEMCPY_HOST_TO_DEVICE);
        aclrtMemcpy((char *)m_input_buffer + 2 * single_chn_size, single_chn_size, m_chw_wrappers[2].data, single_chn_size, ACL_MEMCPY_HOST_TO_DEVICE);
    }

    ret = aclmdlExecute(m_model_id, m_input_dataset, m_output_dataset);

    float scale = std::min((float)m_InputSize.width / (float)img.cols, (float)m_InputSize.height / (float)img.rows);
    void *obuf = nullptr;
    if (m_acl_run_mode == ACL_HOST)
    {
        size_t cp_size = m_output_dims.dims[0] * m_output_dims.dims[1] * m_output_dims.dims[2] * sizeof(float);
        aclrtMemcpy(m_output_buffer_host, cp_size, m_output_buffer, cp_size, ACL_MEMCPY_DEVICE_TO_HOST);
        obuf = m_output_buffer_host;
    }
    else
    {
        obuf = m_output_buffer;
    }
    decode_outputs((float *)obuf, DetObjs, scale, img.cols, img.rows);
    runNms(DetObjs, 0.45);
}

void SampleDetector::runNms(std::vector<BoxInfo> &objects, const float threshold)
{
    auto cmp_lammda = [](const BoxInfo &b1, const BoxInfo &b2)
    { return b1.score < b2.score; };
    std::sort(objects.begin(), objects.end(), cmp_lammda);
    for (int i = 0; i < objects.size(); ++i)
    {
        if (objects[i].score < 0.1)
        {
            continue;
        }
        for (int j = i + 1; j < objects.size(); ++j)
        {
            cv::Rect rect1 = cv::Rect{(int)objects[i].x1, (int)objects[i].y1, int(objects[i].x2 - objects[i].x1), int(objects[i].y2 - objects[i].y1)};
            cv::Rect rect2 = cv::Rect{(int)objects[j].x1, (int)objects[j].y1, int(objects[j].x2 - objects[j].x1), int(objects[j].y2 - objects[j].y1)};
            if (IOU(rect1, rect2) > threshold)
            {
                objects[i].score = 0.f;
            }
        }
    }
    auto iter = objects.begin();
    while (iter != objects.end())
    {
        if (iter->score < 0.1)
        {
            iter = objects.erase(iter);
        }
        else
        {
            ++iter;
        }
    }
}

void SampleDetector::decode_outputs(float *prob, std::vector<BoxInfo> &objects, float scale, const int img_w, const int img_h)
{
    for (int i = 0; i < m_iBoxNums; ++i)
    {
        {

            float x = prob[i + 0 * m_iBoxNums];
            float y = prob[i + 1 * m_iBoxNums];
            float w = prob[i + 2 * m_iBoxNums];
            float h = prob[i + 3 * m_iBoxNums];
            x /= scale;
            y /= scale;
            w /= scale;
            h /= scale;

            int max_cls_pos = 0;
            float max_cls_score = 0.0f;
            for (int j = 0; j < m_iClassNums; ++j)
            {
                auto score = prob[i + (4 + j) * m_iBoxNums];
                if (score > max_cls_score)
                {
                    max_cls_score = score;
                    max_cls_pos = j;
                }
            }

            if (max_cls_score > threshold_)
            {

                cv::Rect box{int(x - w / 2), int(y - h / 2), int(w), int(h)};
                box = box & cv::Rect(0, 0, img_w - 1, img_h - 1);
                if (box.area() > 0)
                {
                    BoxInfo box_info = {
                        float(box.x),
                        float(box.y),
                        float(box.x + box.width),
                        float(box.y + box.height),
                        max_cls_score,
                        max_cls_pos};

                    objects.push_back(box_info);
                }
            }
        }
    }
}
